import time
from multiprocessing import Process, Pipe
from SuperProc import SuperProc

class SmartPool(object):
    def __init__(self, proc_num, fold):
        self.proc_num = proc_num
        self.task_fold = fold
        self.thisSide, self.thatSide = zip(*[Pipe() for _ in range(proc_num)])
        self.thisSideHelp, self.thatSideHelp = zip(*[Pipe() for _ in range(proc_num)])
        self.proc_pool = [SuperProc(p, pH, cnt) for p, pH, cnt in zip(self.thatSide, self.thatSideHelp,range(proc_num))]
        for proc in self.proc_pool:
            proc.start()
        # shut down
        for i in range(proc_num):
            self.thatSide[i].close()
            self.thatSideHelp[i].close()

    def __del__(self):
        for i in range(self.proc_num):
            self.thisSide[i].send(-1)    # switch to helper channel
            self.thatSide[i].close()
            self.thatSideHelp[i].close()

    # add an object of some class, initialize it proc_num=64 times, assigning them to proc_num/fold_num=16 python
    # processes
    def add_target(self, name, lam, args_list=None):
        lam_list = None
        if isinstance(lam, list): lam_list = lam
        for j in range(self.proc_num):
            tuple_list_to_be_send = []
            for i in range(self.task_fold):
                name_fold = name + str(i)
                args = None if args_list is None else args_list[i + j*self.task_fold]
                if lam_list is not None: lam = lam_list[i + j*self.task_fold]
                tuple_list_to_be_send.append((name_fold, lam, args))
            self.thisSide[j].send(0)    # switch to helper channel
            self.thisSideHelp[j].send(tuple_list_to_be_send)

    # if there is index, execute one, otherwise execute all
    def exec_target(self, name, dowhat, args_list = None, index_list = None):
        if index_list is None:
            for j in range(self.proc_num):
                tuple_list_to_be_send = []
                for i in range(self.task_fold):
                    name_fold = name + str(i)
                    args = None if args_list is None else args_list[i + j*self.task_fold]
                    tuple_list_to_be_send.append((name_fold, dowhat, args))
                self.thisSide[j].send(tuple_list_to_be_send)
            res_sort = []
            for j in range(self.proc_num):
                res_sort.extend(self.thisSide[j].recv())
            return res_sort
        else:
            tuple_List_List = [ [None]*self.task_fold ]*self.proc_num
            # sort args
            for index in index_list:
                which_proc = index // self.task_fold
                which_fold = index % self.task_fold
                name_fold = name + str(which_fold)
                args = None if args_list is None else args_list[i + j*self.task_fold]
                tuple_List_List[which_proc][which_fold] = (name_fold, dowhat, args)
                
            # send args
            for j in range(self.proc_num):
                tuple_list_to_be_send = tuple_List_List[j]
                if len(tuple_list_to_be_send) > 0: 
                    self.thisSide[j].send(tuple_list_to_be_send)
                    tuple_List_List[j] = True   # will be return on this channel
                else:
                    tuple_List_List[j] = False

            # receive returns
            for j in range(self.proc_num):
                if not tuple_List_List[j]:
                    continue
                tuple_List_List[j] = self.thisSide[j].recv()

            # sort returns
            res_sort = [None] * len(index_list)
            for i, index in enumerate(index_list):
                which_proc = index // self.task_fold
                which_fold = index % self.task_fold
                res_sort[i] = tuple_List_List[which_proc][which_fold]
            return res_sort

